from chapter04.mnist_dataset import MNISTLoader
from keras.layers import Input, Dense
from keras.models import Model
from keras.optimizers import Adam
from keras.losses import mean_squared_error


def create_sequential_model():
    """
    使用 Functional API 构建一个顺序模型
    :return: model
    """
    input_layer = Input(shape=(28 * 28,), name="input_layer")

    hidden = Dense(64, activation="relu", name="hidden_layer")
    hidden_layer = hidden(input_layer)

    output = Dense(10, activation="softmax", name="output_layer")
    output_layer = output(hidden_layer)

    model = Model(inputs=input_layer, outputs=output_layer, name="Functional_API_Model")
    model.compile(optimizer=Adam(), loss=mean_squared_error, metrics=['accuracy'])
    return model


if __name__ == '__main__':
    dataset = MNISTLoader()

    train_data = dataset.get_train_data()
    train_label = dataset.get_train_label()

    test_data = dataset.get_test_data()
    test_label = dataset.get_test_label()

    model = create_sequential_model()
    print(model.summary())

    x = train_data.reshape((60000, 28 * 28)).astype('float32') / 255
    y = train_label.astype('float32')

    model.fit(x=x, y=y, epochs=10, verbose='auto', batch_size=64)

    x_test = test_data.reshape((10000, 28 * 28)).astype('float32') / 255
    y_test = test_label.astype('float32')

    loss, accuracy = model.evaluate(x=x_test, y=y_test,batch_size=32, verbose=1)

    print(f"Test loss: {loss}")
    print(f"Test accuracy: {accuracy}")

